# -*- coding: utf-8 -*-

# 导入pyspark
from pyspark import SparkContext
from pyspark.mllib.classification import NaiveBayes
from pyspark.mllib.regression import LabeledPoint
from pyspark.mllib.linalg import Vectors
from pyspark.mllib.tree import DecisionTree
import numpy as np

# 导入可视化
import matplotlib
matplotlib.use('Agg') # 不回显
import matplotlib.pyplot as plt


# main函数部分
sc = SparkContext("yarn-client", "NaiveBayes Spark App")

# 加载数据集 
dataPath = "hdfs://192.168.119.141:9100/data/classification/mat_classify"
# 原始数据集
raw_data = sc.textFile(dataPath)
# LabeledPoint数据集
records = raw_data.map(lambda x: x.split(","))
labeledRecords = records.map(lambda x: LabeledPoint(float(x[0]), Vectors.dense([float(y) for y in x[1].split(" ")])))

# 拆分训练集与测试集
data_with_idx = labeledRecords.zipWithIndex().map(lambda (k, v): (v, k))
testRecords = data_with_idx.sample(False, 0.2, 42)
trainRecords = data_with_idx.subtractByKey(testRecords)
testData = testRecords.map(lambda (idx, p): p)
trainData = trainRecords.map(lambda (idx, p): p)
trainData.cache()
testData.cache()

# 写入文件
f = open('getComparisonOfNaiveBayesAndDecisionTree.txt','w')
f.write("\n<br>===LabeledPoint===<br>\n")
f.write(str(labeledRecords.take(5)))
f.write("\n<br>===trainData===<br>\n")
f.write(str(trainData.take(5)))
f.write("\n<br>===testData===<br>\n")
f.write(str(testData.take(5)))
f.write("\n<br>======<br>\n")

# 训练朴素贝叶斯分类模型
model = NaiveBayes.train(trainData, 1.0)
predictionAndLabel = testData.map(lambda p: (model.predict(p.features), p.label))
accuracy = 1.0 * predictionAndLabel.filter(lambda (x, v): x == v).count() / testData.count()

f.write("\n<br>===朴素贝叶斯分类模型===<br>\n")
f.write("\n<br>===(preds, actual)===<br>\n")
f.write(str(predictionAndLabel.take(40)))
f.write("\n<br>===accuracy===<br>\n")
f.write(str(accuracy))
f.write("\n<br>======<br>\n")


# 训练决策树分类模型
modelDecisionTree = DecisionTree.trainClassifier(trainData, 3, {}, maxDepth = 4, maxBins = 8)
predictDecisionTree = modelDecisionTree.predict(testData.map(lambda p: p.features))
predictionAndLabelDecisionTree = testData.map(lambda p: p.label).zip(predictDecisionTree).map(lambda (x, y): (y, x))
accuracyDecisionTree = 1.0 * predictionAndLabelDecisionTree.filter(lambda (x, v): x == v).count() / testData.count()

f.write("\n<br>===决策树分类模型===<br>\n")
f.write("\n<br>===(preds, actual)===<br>\n")
f.write(str(predictionAndLabelDecisionTree.take(40)))
f.write("\n<br>===accuracy===<br>\n")
f.write(str(accuracyDecisionTree))
f.write("\n<br>======<br>\n")

# 关闭文件指针f
f.close()
# 关闭sc
sc.stop()
